#
# s2sdft.py
#
# Copyright (C) 2013 Robert Buj Gelonch
# Copyright (C) 2013 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from operator import itemgetter
import sys

from numpy import abs
from numpy import angle
from numpy import cos
from numpy import int16
from numpy import sin
from numpy import zeros
from numpy.fft import irfft
from numpy.fft import rfft
from pylab import figure
from pylab import plot
from pylab import show
from pylab import subplot
from scipy.io.wavfile import read
from scipy.io.wavfile import write

def module_phase(spectrum):
    return abs(spectrum), angle(spectrum)

def s2sDFT(f1, f2):
    print "s2sDFT.py"
    print
    #--------------------------------------------------------------
    # Read data and obtain basic attributes
    #--------------------------------------------------------------
    [sample_rate_x, x] = read(f1)
    [sample_rate_y, y] = read(f2)
    #--------------------------------------------------------------
    # Obtain the module and the phase
    #--------------------------------------------------------------
    m_x, p_x = module_phase(rfft(x[1:5 * 8000]))
    m_y, p_y = module_phase(rfft(y[1:5 * 8000]))
    #--------------------------------------------------------------
    # Sort the module and the phase
    #--------------------------------------------------------------
    fig = figure(num=None, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')
    
    dm_x = sorted([(k, v) for v, k in zip(m_x, range(len(m_x)))], key=itemgetter(1), reverse=True)
    ax = subplot(211)
    plot([v for k, v in dm_x], 'r', label="|X|")

    dm_y = sorted([(k, v) for v, k in zip(m_y, range(len(m_y)))], key=itemgetter(1), reverse=True)
    plot([v for k, v in dm_y], 'b', label="|Y|")

    dp_x = sorted([(k, v) for v, k in zip(p_x, range(len(p_x)))], key=itemgetter(1), reverse=True)
    ax = subplot(212)
    plot([v for k, v in dp_x], 'r', label="O(X)")

    dp_y = sorted([(k, v) for v, k in zip(p_y, range(len(p_y)))], key=itemgetter(1), reverse=True)
    plot([v for k, v in dp_y], 'b', label="O(Y)")

    show()
    #--------------------------------------------------------------
    # Secret
    #--------------------------------------------------------------
    km_x = [k for k, v in dm_x]
    kp_x = [k for k, v in dp_x]
    #--------------------------------------------------------------
    # Reconstruction: H|X| & O(X)
    #--------------------------------------------------------------
    fig = figure(num=None, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')

    ax = subplot(211)
    rm_x = zeros(len(dm_y))
    for i in range(len(rm_x)):
        rm_x[km_x[i]] = dm_y[i][1]
    plot(rm_x, 'r', label="r(|X|)")
    plot(m_x, 'b', label="|X|")

    ax = subplot(212)
    rp_x = zeros(len(dp_y))
    for i in range(len(rp_x)):
        rp_x[kp_x[i]] = dp_y[i][1]
    plot(rp_x, 'r', label="r(O(X))")
    plot(p_x, 'b', label="O(X)")

    show()
    #--------------------------------------------------------------
    # Signal reconstruction: IFFT
    #--------------------------------------------------------------
    fig = figure(num=None, figsize=(12, 9), dpi=80, facecolor='w', edgecolor='k')

    rX = [complex(rm_x[idx] * cos(rp_x[idx]), rm_x[idx] * sin(rp_x[idx])) for idx in range(len(km_x))]
    rx = irfft(rX)
    plot(rx, 'r', label="reconstruction")
    plot(x, 'b', label="original")
    
    show()
    #--------------------------------------------------------------
    # Save wave file
    #--------------------------------------------------------------
    write("../s2sDFT.R.wav", sample_rate_y, rx.astype(int16))


def main(argv):
    f1 = '../media/fem1.wav'
    f2 = '../media/mal1.wav'
    s2sDFT(f1, f2)

if __name__ == "__main__":
    main(sys.argv[1:])
    print "Done"